#include "pnp_solver.h"
#include <iostream>

PnPSolver::PnPSolver() {
    // 小装甲板3D坐标 (单位: mm)
    small_armor_points_ = {
        cv::Point3f(-67.5, -30.0, 0),   // 左灯条上
        cv::Point3f(-67.5, 0.0, 0),     // 左灯条中  
        cv::Point3f(-67.5, 30.0, 0),    // 左灯条下
        cv::Point3f(67.5, -30.0, 0),    // 右灯条上
        cv::Point3f(67.5, 0.0, 0),      // 右灯条中
        cv::Point3f(67.5, 30.0, 0)      // 右灯条下
    };
    
    // 大装甲板3D坐标
    large_armor_points_ = {
        cv::Point3f(-112.5, -30.0, 0),  // 左灯条上
        cv::Point3f(-112.5, 0.0, 0),    // 左灯条中
        cv::Point3f(-112.5, 30.0, 0),   // 左灯条下
        cv::Point3f(112.5, -30.0, 0),   // 右灯条上
        cv::Point3f(112.5, 0.0, 0),     // 右灯条中
        cv::Point3f(112.5, 30.0, 0)     // 右灯条下
    };
}

bool PnPSolver::loadCameraParams(const std::string& config_path) {
    cv::FileStorage fs(config_path, cv::FileStorage::READ);
    if (!fs.isOpened()) {
        std::cout << "Failed to open camera config file!" << std::endl;
        return false;
    }
    
    fs["camera_matrix"] >> camera_matrix_;
    fs["dist_coeffs"] >> dist_coeffs_;
    
    std::cout << "Camera matrix loaded:" << std::endl;
    std::cout << camera_matrix_ << std::endl;
    
    fs.release();
    return true;
}

bool PnPSolver::solvePnP(Armor& armor) {
    std::vector<cv::Point2f> image_points;
    
    // 这里需要根据实际检测到的灯条点来设置image_points
    // 简化版本：使用装甲板的四个角点
    auto rect_points = armor.getPoints();
    if (rect_points.size() < 4) {
        return false;
    }
    
    // ensure camera params loaded
    if (camera_matrix_.empty()) return false;

    // 选择使用小装甲板还是大装甲板模型
    auto world_points = (armor.getNumber() >= 5) ? large_armor_points_ : small_armor_points_;
    
    // 简化的2D-3D点对应关系（实际需要更精确的匹配）
    image_points = {
        rect_points[0], // 左上
        rect_points[1], // 左下  
        rect_points[2], // 右下
        rect_points[3]  // 右上
    };
    
    // 调整3D点以匹配2D点顺序
    std::vector<cv::Point3f> adjusted_world_points = {
        world_points[0], // 左上
        world_points[2], // 左下
        world_points[5], // 右下
        world_points[3]  // 右上
    };
    
    if (adjusted_world_points.size() != image_points.size()) {
        return false;
    }

    cv::Mat rvec, tvec;
    bool success = false;
    try {
        success = cv::solvePnP(adjusted_world_points, image_points, 
                               camera_matrix_, dist_coeffs_, 
                               rvec, tvec, false, cv::SOLVEPNP_ITERATIVE);
    } catch (const cv::Exception& e) {
        std::cout << "solvePnP exception: " << e.what() << std::endl;
        return false;
    }
    
    if (success) {
        armor.setPose(rvec, tvec);
        armor.setWorldPoints(adjusted_world_points);
        return true;
    }
    
    return false;
}